# -*- coding: utf-8 -*-
"""
Created on Mon Jun 20 11:00:07 2022

@author: Administrator
"""


import cmath
import math
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties
font_set = FontProperties(fname=r"C:\Windows\Fonts\simsun.ttc", size=12)

#求F范数
def Fro(FA)    :
     M= np.size(FA,0)
     N= np.size(FA,1)
     sum1=0
     for i in range(M):
         for j in range(N):
             sum1+=FA[i,j]**2
     return sum1

def sparse_learning(Aa_old,Y,iters):
    MIN_GAMMA       = 1e-16
  




#*** Initializations ***初始化
    M= np.size(Aa_old,0) #L是网格数
    L=np.size(Aa_old,1)
    T = np.size(Y,1)#T是快拍数,M是阵元
    gamma = np.ones((L))
    C=np.zeros((M-1),dtype=complex)
  #初始化超参数向量
 

    sigma_2=1
    k = 0
    a=1e-4
    b=1e-4
    p=0.01

    while k<iters:#iters
        Gamma=np.diag(gamma)#L*L维度
        D=np.array(np.append(0,C.T)).T
        PS=np.diag(D)
        Aa_1=Aa_old+np.dot(PS,Aa_old)
        Sigma_x=sigma_2*np.eye(M)+np.dot(np.dot(Aa_1,Gamma),np.conj(Aa_1).T)
        Sigma_s=Gamma-np.dot(np.dot(np.dot(np.dot(Gamma,np.conj(Aa_1).T), np.linalg.inv(Sigma_x)),Aa_1),Gamma)
        M_u=sigma_2**(-1)*np.dot(np.dot(Sigma_s,np.conj(Aa_1).T),Y)
        GP_1=np.zeros((M,M))
        GP_2=np.zeros((M,M))
        Q_t=np.zeros((M-1,M-1))
        Q=np.zeros((M-1))
        Q_A=np.zeros((M-1))
        S_S=(np.dot(M_u,np.conj(M_u).T)+T*Sigma_s)
        for f in range(M-1):
            for h in range(M-1):
                GP_1[f+1,f+1]=1
                GP_2[h+1,h+1]=1
                Q_t[f,h]=np.trace(np.dot(np.dot(np.dot(np.dot(np.conj(GP_1).T,GP_2),Aa_old),S_S),np.conj(Aa_old).T))
                GP_1=np.zeros((M,M))
                GP_2=np.zeros((M,M))
    #QtH*X
        for r in range(M-1):
            GP_1[r+1,r+1]=1
            Q[r]=np.trace(np.dot(np.dot(np.dot(np.conj(GP_1).T,Y),np.conj(M_u).T),np.conj(Aa_old).T))
            GP_1=np.zeros((M,M))
        for u in range(M-1):
            GP_1[u+1,u+1]=1
            Q_A[u]=np.trace(np.dot(np.dot(np.dot(np.conj(GP_1).T,Aa_old),S_S),np.conj(Aa_old).T))
            GP_1=np.zeros((M,M))
        C=np.dot(np.linalg.inv(Q_t),(Q-Q_A))#误差系数
        FA=Y-np.dot(Aa_1,M_u)
        sigma_2=(M*T+a-1)**(-1)*(b+Fro(FA)+T*np.trace(np.dot(np.dot(Aa_1,Sigma_s),np.conj(Aa_1).T)))#噪声方差
        gamma_old = gamma.copy()
        gamma_old=gamma.copy()
        m_u=np.array(M_u.copy())**2
        gamma=np.sum(abs(m_u),axis=1)+np.diag(Sigma_s)
        gamma=abs((-T+(T**2+4*p*gamma)**(0.5))/(2*p))
        #迭代条件
        if np.linalg.norm(gamma-gamma_old,2)/np.linalg.norm(gamma_old,2)<=MIN_GAMMA:
            break
        else:
            k=k+1
    return gamma
derad =math.pi/180      #角度->弧度
N = 8              # 阵元个数        
T=200              #快拍数
theta =[-40,0,20,60]  # 待估计角度，
M =len(theta)              # 信源数目
snr = 10            # 信噪比
grid_interval=2  # 网格间距
Azimuth_grid=np.arange(-90,90,grid_interval)
dd = 0.5           # 阵元间距 
si=-0.5+np.random.rand(1,N)
rho=np.zeros((1,N),dtype=complex)
si[0][0]=0
for s in range(N):
     rho[:,s]=(1+np.sqrt(12)*0.25*si[:,s])*np.exp(1j*np.sqrt(12)*35*derad*si[:,s]*derad)-1#幅相误差减1
T_xishu=np.diag(rho[0].T)
d=np.arange(0,N*dd,dd)
A=np.zeros((N,M),dtype=complex)
for i in range(N):
    for n in range(M):
        A[i,n]=cmath.exp(1j*math.pi*(i-1)*math.sin(theta[n]*derad))
Az=np.zeros((len(Azimuth_grid)))
Aa=np.ones((N,len(Azimuth_grid)),dtype=complex)
for n in range(N):
    for ang in range(len(Azimuth_grid)):
        Az[ang]=Azimuth_grid[ang]*derad;
        Aa[n,ang]=np.exp(1j*np.pi*(n-1)*np.sin(Az[ang]))
Aa_old=Aa.copy()
A=A+np.dot(T_xishu,A)
S=(np.random.randn(M,T)+1j*np.random.randn(M,T))
Vj=np.diag(np.sqrt(10**(snr/10)/np.diag(1/T*(np.dot(S,np.conj(S).T)))))
S=np.dot(Vj,S) # 产生信号
noise=np.sqrt(1/2)*(np.random.randn(N,T)+1j*np.random.randn(N,T))# 加噪声
Y=np.dot(A,S)+noise
gamma=sparse_learning(Aa_old,Y,1000)
plt.figure()
plt.plot(Azimuth_grid,gamma, label = 'SBL')
plt.xlim(-90,90)
plt.xlabel('入射角/(degree)',fontproperties=font_set)
plt.ylabel('空间谱/(dB)',fontproperties=font_set)
plt.grid()
plt.show()
